import json
import re

file_name = 'a' #'demo'

with open(file_name+'.trace.json', 'r') as f:
    trace_json_txt = f.read()

#~ data = json.load ( trace_json_txt )
data = eval ( trace_json_txt )

code = data['code']
linecodes = code.split('\n')
trace = data['trace']

print ( '\n### TRACE: ', json.dumps(trace, sort_keys=True, indent=2) )
#~ print ( 'CODE: ',  json.dumps(code, sort_keys=True, indent=2) )
print ( '\n### CODE:'); print( code )

######### plain trace ########
for frame in trace:
    if frame['event'] == 'step_line':
        lineno = int( frame['line'] )
        print ( lineno, linecodes[lineno-1] )


def not_empty( line ):
    return line.strip(' \t') and not line.strip(' \t').startswith("#") #is not empty and not comment

########### Construct LOOPS Tree info #######


#~ LOOP_INDENTS = [0] * len(linecodes)
#~ LOOP_DEPTHS  = [0] * len(linecodes)
#~ LOOP_STARTS  = [None] * len(linecodes)

LOOP_INDENTS = []  # for each line
LOOP_DEPTHS  = []  # we'll store its
LOOP_STARTS  = []  # loop dependency

def construct_loops_info():
    loop_depth = 0
    LOOPS_STACK = [  {'indent':0, 'start':'-'} ]   # for tree of loops: (indent, start) 
    GET_INSIDELOOP_INDENT = False

    for nr, line in enumerate(linecodes):
        line = line.replace('\t', '    ')  # replace tabs
        if not_empty(line):  #if line not empty
        
            indent = len(line) - len( line.lstrip(' ') )

            while indent < LOOPS_STACK[-1]['indent']:  # if we return from some loop
                LOOPS_STACK.pop()   
                loop_depth -= 1
                GET_INSIDELOOP_INDENT = False

            if GET_INSIDELOOP_INDENT:     # if we are first line inside the loop -- get indentation
                LOOP_INDENTS[-1] = LOOPS_STACK[-1]['indent'] = indent
                GET_INSIDELOOP_INDENT = False
                
            linecode = line.lstrip(' ') .replace('\t', ' ').replace('\\', ' ').replace('(', ' ')
            if linecode.startswith("for ") or linecode.startswith("while "): # TODO migth better be regexp (.*?)\W
                loop_depth += 1;
                LOOPS_STACK.append( {'indent': indent+1, 'start': nr} )  # indent+1 is a workaround for our model to deal with oneliner-loop
                GET_INSIDELOOP_INDENT = True

            
        LOOP_INDENTS.append( LOOPS_STACK[-1]['indent'] )
        LOOP_DEPTHS.append( loop_depth )
        LOOP_STARTS.append( LOOPS_STACK[-1]['start'] ) # line number, where loop starts

        if True:
            #~ print("\nDBG construct_loops_info:\n")
            print("%3d"%nr, "INDENT, DEPTH, START:", "%3d %3d %3s"%(LOOP_INDENTS[-1],  LOOP_DEPTHS[-1], LOOP_STARTS[-1]),":", line)

construct_loops_info()

###########  nested TABLE trace #######

from collections import OrderedDict, defaultdict

def order_dicts( d ):
    ordered =  OrderedDict( sorted(d.items(), key=lambda t: t[0]) )
    for key, val in ordered.items():
        if isinstance( val[0], dict ):
            val[0] = order_dicts( val[0] )
    return ordered

from pprint import pprint


GO='+'; SKIP=' '
trace_rows = defaultdict(list)
trace_stack = [trace_rows]


def get_frame_indent(frameno):
    lineno = int(trace[frameno]['line'])
    line =  linecodes[lineno-1]
    indent = len(line) - len( line.lstrip(' ') )
    return indent


        
def repetition_trace(start_frameno):

    start_lineno = int(trace[start_frameno]['line'])
    start_indent = get_frame_indent(start_frameno)
    lineno = last_lineno = start_lineno

    frameno = start_frameno
    looping = False
    
    while frameno  < len(trace):

        frame = trace[frameno]
        
        lineno = int( frame['line'] )
        line =  linecodes[lineno-1]
        linecode = line.lstrip(' ')
        indent = len(line) - len(linecode)

        # check for indentation block end
        if indent < start_indent:
            break
        if looping and indent==start_indent and lineno > start_lineno:  # some idioms not included: for... else...
            break

        # process line
        if frame['event'] == 'step_line':

            if lineno > last_lineno+1:
                for tmp_lineno in range(last_lineno+1, lineno):
                    trace_stack[-1][tmp_lineno].append( SKIP )

            #~ trace_stack[-1][lineno].append( GO )
            def get_info():
                import ast
                def get_assignment_target(code):  # http://hg.python.org/cpython/file/3.3/Lib/ast.py#l89
                    """
                    >>> print( get_assignment_target("pinigai = laikas * atlyginimas") )
                    pinigai
                    """
                    node = ast.parse(code, mode='exec')
                    if (isinstance(node.body[0], ast.Assign)):
                        assignment = node.body[0]
                        return assignment.targets[0].id
                # http://eli.thegreenplace.net/2009/11/28/python-internals-working-with-python-asts/ might be of use
                #~ global linecode
                nonlocal linecode
                if '#' in linecode:
                    linecode, comment = linecode.split('#', 1)
                    linecode = linecode.strip()
                    
                code_chunks = re.split('(\W+)', linecode)
                code_chunks = [x.strip() for x in code_chunks]
                start_chunk = code_chunks[0]
                #~ print ('start_chunk', start_chunk)
                
                if code_chunks[0] in ['if', 'elif', 'while']:
                    expr = linecode.split(':')[0]
                    expr = expr[len(start_chunk ):].strip()

                elif code_chunks[0] in ['for']: 
                    #~ expr = code_chunks[1]
                    expr = re.split('\sin[\s\[]', linecode)[0]  # TODO: better use ast here...   in case in is separated by 'for x in[3, 5]' or other weird cases
                    expr = expr[len(start_chunk ):].strip()

                elif linecode.endswith(":"):
                    expr = ''
                    
                elif linecode.startswith("print"):  # optionally could track what is printed
                    expr = '' 

                elif linecode == '':
                    expr = None
                    
                else:
                    expr = get_assignment_target(linecode) or '' # target variable name or None-->''

                if expr:
                    value = eval(expr, trace[frameno+1]['globals'])  # use globals from next trace -- as the newly assigned vaues will be there
                else:
                    value = '' if expr is None else GO
                                        
                #~ return "%s@%s" % (indent, frameno)
                #~ return "%s@%s:%s=%s" % (indent, frameno, expr, value)
                return value
                   
            trace_stack[-1][lineno].append( get_info() )            
            last_lineno = lineno   

            if linecode.startswith('for') \
            or linecode.startswith('while'):
                if lineno > start_lineno:
                    new_loop_trace = defaultdict(list)
                    trace_stack[-1][lineno][-1] = new_loop_trace  # replace GO
                    trace_stack.append( new_loop_trace )
                    frameno = repetition_trace(frameno)   # give frameno, and update it after return ;)
                    sub_trace = trace_stack.pop()
                    last_lineno = max( sub_trace.keys() )
                    #~ print ( "\n Finished branch @ %s..%s " % (lineno, last_lineno) )
                    #~ pprint ( order_dicts( sub_trace ) )
                else:
                    looping = True

        frameno += 1

    return frameno-1
                

         
print("\n    Result  \n")
repetition_trace(0)


########### TEST case ############

test_trace_rows  = {
1: [GO],
3: [ {
    3: [GO,GO,GO,GO,GO,] ,
    4: [GO,GO,GO,GO,GO,] ,
    5: [GO,GO,GO,GO,GO,] ,
    6: [SKIP,SKIP,GO,SKIP,SKIP,] ,
    7: [GO,GO,GO,GO,GO,] 
 } ],
9: [GO],
}
# http://docs.python.org/3.3/library/collections.html#ordereddict-examples-and-recipes

print("\n    Expected  \n")

pprint ( order_dicts( test_trace_rows ) )
#~ print ( json.dumps( order_dicts( test_trace_rows ), indent=2)  )

"""####################
s=0

for x in range(1,6):
    s=s+x
    if x%3==0:
        s=0
    print(x, s)

print ("finish")
########################

1 s=0
3 for x in range(1,6):
4     s=s+x
5     if x%3==0:
7     print(x, s)
3 for x in range(1,6):
4     s=s+x
5     if x%3==0:
7     print(x, s)
3 for x in range(1,6):
4     s=s+x
5     if x%3==0:
6         s=0
7     print(x, s)
3 for x in range(1,6):
4     s=s+x
5     if x%3==0:
7     print(x, s)
3 for x in range(1,6):
4     s=s+x
5     if x%3==0:
7     print(x, s)
3 for x in range(1,6):
9 print ("finish")

#######################

[
 ['+'],
 [    ['+', '+', '+', '+', '+'],
      ['+', '+', '+', '+', '+'],
      ['+', '+', '+', '+', '+'],
      [' ', ' ', '+', ' ', ' '],
      ['+', '+', '+', '+', '+']
 ],
 ['+']
 ]


#********** 2 level loop demo ***************
s = 5
for x in [1, 3, 2]:
    print('1st', x)
    if x > 2:
        for z in range(2): print('2nd')
        print('bla')
        for i in range(x):
            a = i
            print('3rd', i)
    # print('back to 1st')
print('end')
        
"""
